"""
数据库和文件备份脚本
支持自动备份和定时任务
"""
import os
import sys
import subprocess
import shutil
from datetime import datetime
from pathlib import Path

# 添加父目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.core.config import settings
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def backup_database():
    """
    备份 PostgreSQL 数据库
    使用 pg_dump 命令
    """
    try:
        # 生成备份文件名
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        backup_dir = Path("./backups/database")
        backup_dir.mkdir(parents=True, exist_ok=True)
        
        backup_file = backup_dir / f"backup_{timestamp}.sql"
        
        logger.info(f"开始备份数据库...")
        logger.info(f"备份文件: {backup_file}")
        
        # 构建 pg_dump 命令
        # 注意：需要设置 PGPASSWORD 环境变量或使用 .pgpass 文件
        env = os.environ.copy()
        env['PGPASSWORD'] = settings.DATABASE_PASSWORD
        
        cmd = [
            'pg_dump',
            '-h', settings.DATABASE_HOST,
            '-p', str(settings.DATABASE_PORT),
            '-U', settings.DATABASE_USER,
            '-d', settings.DATABASE_NAME,
            '-F', 'p',  # plain text format
            '-f', str(backup_file)
        ]
        
        # 执行备份
        result = subprocess.run(
            cmd,
            env=env,
            capture_output=True,
            text=True
        )
        
        if result.returncode == 0:
            # 获取文件大小
            file_size = backup_file.stat().st_size / (1024 * 1024)  # MB
            logger.info(f"✅ 数据库备份成功！")
            logger.info(f"   文件大小: {file_size:.2f} MB")
            
            # 清理旧备份（保留最近 7 天）
            cleanup_old_backups(backup_dir, days=7)
            
            return True
        else:
            logger.error(f"❌ 数据库备份失败！")
            logger.error(f"   错误信息: {result.stderr}")
            return False
            
    except FileNotFoundError:
        logger.error("❌ 未找到 pg_dump 命令！")
        logger.error("   请确保 PostgreSQL 客户端工具已安装并添加到 PATH")
        return False
    except Exception as e:
        logger.error(f"❌ 备份过程出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def backup_files():
    """
    备份上传的文件和向量数据
    """
    try:
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        backup_dir = Path("./backups/files")
        backup_dir.mkdir(parents=True, exist_ok=True)
        
        backup_file = backup_dir / f"files_{timestamp}.tar.gz"
        
        logger.info(f"开始备份文件...")
        logger.info(f"备份文件: {backup_file}")
        
        # 要备份的目录
        dirs_to_backup = [
            './data/uploads',
            './data/knowledge_base',
            './data/vector_store',
        ]
        
        # 使用 tar 命令压缩
        import tarfile
        
        with tarfile.open(backup_file, 'w:gz') as tar:
            for dir_path in dirs_to_backup:
                if os.path.exists(dir_path):
                    tar.add(dir_path, arcname=os.path.basename(dir_path))
                    logger.info(f"   已添加: {dir_path}")
        
        # 获取文件大小
        file_size = backup_file.stat().st_size / (1024 * 1024)  # MB
        logger.info(f"✅ 文件备份成功！")
        logger.info(f"   文件大小: {file_size:.2f} MB")
        
        # 清理旧备份
        cleanup_old_backups(backup_dir, days=7)
        
        return True
        
    except Exception as e:
        logger.error(f"❌ 文件备份失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def cleanup_old_backups(backup_dir: Path, days: int = 7):
    """
    清理旧备份文件
    
    Args:
        backup_dir: 备份目录
        days: 保留天数
    """
    try:
        from datetime import timedelta
        
        cutoff_time = datetime.now() - timedelta(days=days)
        
        deleted_count = 0
        for backup_file in backup_dir.glob('*'):
            if backup_file.is_file():
                file_time = datetime.fromtimestamp(backup_file.stat().st_mtime)
                if file_time < cutoff_time:
                    backup_file.unlink()
                    deleted_count += 1
                    logger.info(f"   已删除旧备份: {backup_file.name}")
        
        if deleted_count > 0:
            logger.info(f"   清理了 {deleted_count} 个旧备份文件")
            
    except Exception as e:
        logger.warning(f"⚠️  清理旧备份失败: {str(e)}")


def restore_database(backup_file: str):
    """
    恢复数据库
    
    Args:
        backup_file: 备份文件路径
    """
    try:
        logger.info(f"开始恢复数据库...")
        logger.info(f"备份文件: {backup_file}")
        
        if not os.path.exists(backup_file):
            logger.error(f"❌ 备份文件不存在: {backup_file}")
            return False
        
        # 警告
        logger.warning("⚠️  警告：此操作将覆盖当前数据库！")
        confirm = input("确认恢复？(yes/no): ").strip().lower()
        if confirm != 'yes':
            logger.info("❌ 取消恢复操作")
            return False
        
        # 构建 psql 命令
        env = os.environ.copy()
        env['PGPASSWORD'] = settings.DATABASE_PASSWORD
        
        cmd = [
            'psql',
            '-h', settings.DATABASE_HOST,
            '-p', str(settings.DATABASE_PORT),
            '-U', settings.DATABASE_USER,
            '-d', settings.DATABASE_NAME,
            '-f', backup_file
        ]
        
        # 执行恢复
        result = subprocess.run(
            cmd,
            env=env,
            capture_output=True,
            text=True
        )
        
        if result.returncode == 0:
            logger.info(f"✅ 数据库恢复成功！")
            return True
        else:
            logger.error(f"❌ 数据库恢复失败！")
            logger.error(f"   错误信息: {result.stderr}")
            return False
            
    except Exception as e:
        logger.error(f"❌ 恢复过程出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """主函数"""
    logger.info("=" * 60)
    logger.info("思和智库 - 备份工具")
    logger.info("=" * 60)
    logger.info()
    
    # 备份数据库
    db_success = backup_database()
    logger.info()
    
    # 备份文件
    files_success = backup_files()
    logger.info()
    
    # 总结
    logger.info("=" * 60)
    if db_success and files_success:
        logger.info("✅ 备份完成！")
    else:
        logger.warning("⚠️  备份部分失败，请检查日志")
    logger.info("=" * 60)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='数据库备份和恢复工具')
    parser.add_argument('--restore', type=str, help='恢复数据库（指定备份文件路径）')
    
    args = parser.parse_args()
    
    if args.restore:
        restore_database(args.restore)
    else:
        main()

